#! /usr/bin/python3
# -*- coding: UTF-8 -*-
  
'''※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※
File Name: Image_Dataset.py
Author: GID5564
Description: 划分数据集与测试集
Version: 1.0
Created Time: 23/05/24-05:23:05
※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※'''

import os
import shutil
import random
import pandas as pd



class ImageDataset():
    def __init__(self,dataset_path="dataset",frac=0.2):
        #指定数据集路径
        self.dataset_path=dataset_path
        self.frac=frac
        self.classes=os.listdir(self.dataset_path)
        self.df=pd.DataFrame()
        
        #在 train 和 test 文件夹中创建各类别子文件夹
        self.__create_dir('test')
        self.__create_dir('train')


    def __create_dir(self,dir_name):
        data_dir=os.path.join(self.dataset_path, dir_name)
        if not os.path.exists(data_dir):
            # 创建  文件夹
            os.mkdir(data_dir)
        # 在 train 和 test 文件夹中创建各类别子文件夹
        for var in self.classes:
            part_data_dir=os.path.join(data_dir, var)
            if not os.path.exists(part_data_dir):
                os.mkdir(part_data_dir)
            
    def __move_image(self,images,dir_name,var):
        for image in images:
            old_img_path = os.path.join(self.dataset_path, str(var), image) # 获取原始文件路径
        
            new_path = os.path.join(self.dataset_path, dir_name, str(var), image) # 获取 test 目录的新文件路径

            shutil.move(old_img_path, new_path) # 移动文件
 
    def split(self):
        random.seed(123) # 随机数种子，便于复现
        for var in self.classes:
            # 读取该类别的所有图像文件名
            old_dir = os.path.join(self.dataset_path, var)
            images_filename = os.listdir(old_dir)
            random.shuffle(images_filename) # 随机打乱
        
            # 划分训练集和测试集
            testset_numer = int(len(images_filename) * self.frac) # 测试集图像个数
            testset_images = images_filename[:testset_numer]      # 获取拟移动至 test 目录的测试集图像文件名
            trainset_images = images_filename[testset_numer:]     # 获取拟移动至 train 目录的训练集图像文件名


            # 移动图像至 test 目录
            self.__move_image(testset_images,'test',var)

            # 移动图像至 train 目录
            self.__move_image(trainset_images,'train',var)


            # 删除旧文件夹
            assert len(os.listdir(old_dir)) == 0 # 确保旧文件夹中的所有图像都被移动走
            shutil.rmtree(old_dir) # 删除文件夹
            
            # 保存到表格中
            self.df = self.df.append({'class':var, 'trainset':len(trainset_images), 'testset':len(testset_images)}, ignore_index=True)
            
    def tocsv(self):
        df=self.df
        # 数据集各类别数量统计表格，导出为 csv 文件
        df['total'] = df['trainset'] + df['testset']
        df.to_csv('数据量统计.csv', index=False)

        


ids=ImageDataset()
ids.split()
ids.tocsv()




























